. ./path.sh || exit 1;
export CUDA_VISIBLE_DEVICES="0,1"
export NCCL_DEBUG=INFO
export CUDA_LAUNCH_BLOCKING=1
stage=4 # start from 0 if you need to start from data preparation
stop_stage=4
num_nodes=1
node_rank=0

data=/home/backup_nfs2/nfs1_data/data_aishell
data_url=www.openslr.org/resources/33

dict=./output/dict/aishell_gxl_data/lang_char.txt
data_gxl=./output/aishell_gxl_data
exp_gxl=./output/aishell_gxl_exp

nj=16
data_type=raw

train_set=train
dev_set=dev

train_config=conf/tuning/train_asr_whisper_full.yaml
cmvn=false

checkpoint=
num_workers=10
prefetch=0

average_checkpoint=true
decode_checkpoint=$dir/final.pt
average_num=30
decode_modes="ctc_greedy_search ctc_prefix_beam_search attention attention_rescoring"

deepspeed=false

. tools/parse_options.sh || exit 1; # 解析命令传参

if [ ${stage} -le -1 ] && [ ${stop_stage} -ge -1 ]; then
  echo "stage -1: Data Download"
  local/download_and_untar.sh ${data} ${data_url} data_aishell
  local/download_and_untar.sh ${data} ${data_url} resource_aishell
fi

if [ ${stage} -le 0 ] && [ ${stop_stage} -ge 0 ]; then
  # Data preparation
  echo 'Data preparation'
  local/aishell_data_prep.sh ${data}/wav \
    ${data}/transcript
fi

if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
  # remove the space between the text labels for Mandarin dataset,
  echo "remove the space between the text labels for Mandarin dataset"
  for x in train dev test; do
    cp ${data_gxl}/${x}/text ${data_gxl}/${x}/text.org
    paste -d " " <(cut -f 1 -d" " ${data_gxl}/${x}/text.org) \
      <(cut -f 2- -d" " ${data_gxl}/${x}/text.org | tr -d " ") \
      > ${data_gxl}/${x}/text
   rm ${data_gxl}/${x}/text.org
  done
  if [ ${cmvn} == true ]; then
    # 如果A为true，则执行
    echo "执行cmvn"
    echo "compute cmvn states"
    mkdir -p ${data_gxl}/${train_set}
    tools/compute_cmvn_stats.py --num_workers 16 --train_config $train_config \
      --in_scp ${data_gxl}/${train_set}/wav.scp \
      --out_cmvn ${data_gxl}/$train_set/global_cmvn
  else
      # 如果A为false，则不执行
      echo "变量A为false，不执行cmvn"
  fi
fi

if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
  echo "Make a dictionary"
  mkdir -p $(dirname $dict)
  echo "<blank> 0" > ${dict}  # 0 is for "blank" in CTC
  echo "<unk> 1"  >> ${dict}  # <unk> must be 1
  tools/text2token.py -s 1 -n 1 ${data_gxl}/train/text | cut -f 2- -d" " \
    | tr " " "\n" | sort | uniq | grep -a -v -e '^\s*$' | \
    awk '{print $0 " " NR+1}' >> ${dict}
  num_token=$(cat $dict | wc -l)
  echo "<sos/eos> $num_token" >> $dict
fi

if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
  echo "Prepare data, prepare required format"
  for x in dev test ${train_set}; do
    if [ $data_type == "shard" ]; then
      tools/make_shard_list.py --num_utts_per_shard $num_utts_per_shard \
        --num_threads 16 data/$x/wav.scp data/$x/text \
        $(realpath data/$x/shards) data/$x/data.list
    else
      tools/make_raw_list.py $data_gxl/$x/wav.scp $data_gxl/$x/text \
        $data_gxl/$x/data.list
    fi
  done
fi
if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
  mkdir -p $exp_gxl
  # You have to rm `INIT_FILE` manually when you resume or restart a
  # multi-machine training.
  INIT_FILE=$exp_gxl/ddp_init
  rm -f ${INIT_FILE}  # remove previous INIT_FILE
  init_method=file://$(readlink -f $INIT_FILE)
  echo "$0: init method is $init_method"
  num_gpus=$(echo $CUDA_VISIBLE_DEVICES | awk -F "," '{print NF}')
  # Use "nccl" if it works, otherwise use "gloo"
  # 分布式通信的初始化的后端实现, gpu推荐nccl,cpu推荐loo
  dist_backend="nccl"
  world_size=`expr $num_gpus \* $num_nodes`
  echo "total gpus is: $world_size"
  cmvn_opts=
  $cmvn && cp ${data_gxl}/${train_set}/global_cmvn $exp_gxl
  $cmvn && cmvn_opts="--cmvn ${dir}/global_cmvn"

  # train.py rewrite $train_config to $exp_gxl/train.yaml with model input
  # and output dimension, and $exp_gxl/train.yaml will be used for inference
  # and export.
  echo "using torch ddp"
  for ((i = 0; i < $num_gpus; ++i)); do
  {
    gpu_id=$(echo $CUDA_VISIBLE_DEVICES | cut -d',' -f$[$i+1]) # 返回索引为i+1的值,此处的索引从1开始算
    # Rank of each gpu/process used for knowing whether it is
    # the master of a worker.
    rank=`expr $node_rank \* $num_gpus + $i`
#    python $WENET_DIR/wenet/bin/train1.py
    python $WENET_DIR/wenet/bin/train.py \
      --gpu $gpu_id \
      --config $train_config \
      --data_type $data_type \
      --symbol_table $dict \
      --train_data $data_gxl/$train_set/data.list \
      --cv_data $data_gxl/$dev_set/data.list \
      ${checkpoint:+--checkpoint $checkpoint} \
      --model_dir $exp_gxl \
      --ddp.init_method $init_method \
      --ddp.world_size $world_size \
      --ddp.rank $rank \
      --ddp.dist_backend $dist_backend \
      --num_workers ${num_workers} \
      --prefetch ${prefetch} \
      $cmvn_opts \
      --pin_memory
  } &
  done
  wait
fi